import time
import cv2
import numpy as np
from rknnlite.api import RKNNLite

class Detect:
    def __init__(self, conf):
    # 在运行推理前，需要获取一张图像来初始化缩放填充参数 load_frame_info
        self.rknn= RKNNLite(verbose=False)
        self.name = conf["name"]   
        self.rknn.load_rknn(conf["model"]) #注意改路径需要人为在前面的传递参数前加上工作环境路径
        self.rknn.init_runtime(core_mask = RKNNLite.NPU_CORE_AUTO)
        self.target_h = conf["input"]["height"]
        self.target_w = conf["input"]["width"]
        self.outputs = None
        self.orig_size = None
        self.padding_w = None 
        self.padding_h = None
        #(ratio,pw,ph)
        self.conf_thres=conf["conf_thres"]
        self.iou_thres=conf["iou_thres"]
        self.pad_color = (0,0,0)
        self.scale = 1
        self.tmp_target_size=None

    def predict(self, frame):
        img_input = self.load_frame(frame)
        self.outputs = self.rknn.inference(inputs=[img_input])
        return self.post_process(self.outputs)
    
    def load_frame_info(self,frame):
        # 推理之前必做一次
        shape = frame.shape[:2]  # current shape [height, width]
        self.orig_size = shape
        self.scale = min(self.target_h / shape[0], self.target_w / shape[1])
        self.tmp_target_size = (int(round(shape[0] *self.scale)), int(round(shape[1] * self.scale)))
        dh, dw =self.target_h - self.tmp_target_size[0], self.target_w - self.tmp_target_size[1]  # wh padding

        dh /= 2  # divide padding into 2 sides
        dw /= 2
        self.padding_h = dh
        self.padding_w = dw
        
    def load_frame(self, img):
        if img.shape[:2]!=self.tmp_target_size:
            img = cv2.resize(img, self.tmp_target_size[::-1], interpolation=cv2.INTER_LINEAR) #缩放
        top, bottom = int(round(self.padding_h- 0.1)), int(round(self.padding_h + 0.1))
        left, right = int(round(self.padding_w - 0.1)), int(round(self.padding_w + 0.1))
       
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=self.pad_color)
        return img.reshape(1,*img.shape).astype(np.float32)
        
    def filter_boxes(self,boxes, box_confidences, box_class_probs):
        """Filter boxes with object threshold.
        """
        box_confidences = box_confidences.reshape(-1)
        class_max_score = np.max(box_class_probs, axis=-1)
        classes = np.argmax(box_class_probs, axis=-1)

        _class_pos = np.where(class_max_score* box_confidences >= self.conf_thres)
        scores = (class_max_score* box_confidences)[_class_pos]

        boxes = boxes[_class_pos]
        classes = classes[_class_pos]

        return boxes, classes, scores

    def nms_boxes(self,boxes, scores):
        """Suppress non-maximal boxes.
        # Returns
            keep: ndarray, index of effective boxes.
        """
        x = boxes[:, 0]
        y = boxes[:, 1]
        w = boxes[:, 2] - boxes[:, 0]
        h = boxes[:, 3] - boxes[:, 1]

        areas = w * h
        order = scores.argsort()[::-1]

        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)

            xx1 = np.maximum(x[i], x[order[1:]])
            yy1 = np.maximum(y[i], y[order[1:]])
            xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
            yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

            w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
            h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
            inter = w1 * h1

            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(ovr <= self.iou_thres)[0]
            order = order[inds + 1]
        keep = np.array(keep)
        return keep
   
    @staticmethod
    def dfl(position):
        x = np.array(position)
        n, c, h, w = x.shape
        p_num = 4
        mc = c // p_num
        y = x.reshape(n, p_num, mc, h, w)
        # Softmax along the mc dimension (axis=2)
        y = np.exp(y - np.max(y, axis=2, keepdims=True))
        y = y / np.sum(y, axis=2, keepdims=True)
        acc_metrix = np.arange(mc).astype(np.float32).reshape(1, 1, mc, 1, 1)
        y = (y * acc_metrix).sum(axis=2)
        return y



    def box_process(self,position):
        grid_h, grid_w = position.shape[2:4]
        col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h))
        col = col.reshape(1, 1, grid_h, grid_w)
        row = row.reshape(1, 1, grid_h, grid_w)
        grid = np.concatenate((col, row), axis=1)
        stride = np.array([self.target_h//grid_h, self.target_w//grid_w]).reshape(1,2,1,1)

        position = self.dfl(position)
        box_xy  = grid +0.5 -position[:,0:2,:,:]
        box_xy2 = grid +0.5 +position[:,2:4,:,:]
        xyxy = np.concatenate((box_xy*stride, box_xy2*stride), axis=1)

        return xyxy

    def post_process(self,input_data):
        boxes, scores, classes_conf = [], [], []
        defualt_branch=3
        pair_per_branch = len(input_data)//defualt_branch
        # Python 忽略 score_sum 输出
        for i in range(defualt_branch):
            boxes.append(self.box_process(input_data[pair_per_branch*i]))
            classes_conf.append(input_data[pair_per_branch*i+1])
            scores.append(np.ones_like(input_data[pair_per_branch*i+1][:,:1,:,:], dtype=np.float32))

        def sp_flatten(_in):
            ch = _in.shape[1]
            _in = _in.transpose(0,2,3,1)
            return _in.reshape(-1, ch)

        boxes = [sp_flatten(_v) for _v in boxes]
        classes_conf = [sp_flatten(_v) for _v in classes_conf]
        scores = [sp_flatten(_v) for _v in scores]

        boxes = np.concatenate(boxes)
        classes_conf = np.concatenate(classes_conf)
        scores = np.concatenate(scores)

        # filter according to threshold
        boxes, classes, scores = self.filter_boxes(boxes, scores, classes_conf)

        # nms
        nboxes, nclasses, nscores = [], [], []
        for c in set(classes):
            inds = np.where(classes == c)
            b = boxes[inds]
            c = classes[inds]
            s = scores[inds]
            keep = self.nms_boxes(b, s)

            if len(keep) != 0:
                nboxes.append(b[keep])
                nclasses.append(c[keep])
                nscores.append(s[keep])

        if not nclasses and not nscores:
            return None, None, None

        boxes = np.concatenate(nboxes)
        classes = np.concatenate(nclasses)
        scores = np.concatenate(nscores)

        return boxes, classes, scores